package com.pig4cloud.pigx.knowledge.support.rule;

import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.pig4cloud.pigx.knowledge.dto.AiMessageResultDTO;
import com.pig4cloud.pigx.knowledge.dto.ChatMessageDTO;
import com.pig4cloud.pigx.knowledge.dto.SiliconflowImageModelDTO;
import com.pig4cloud.pigx.knowledge.entity.AiModelEntity;
import com.pig4cloud.pigx.knowledge.mapper.AiModelMapper;
import com.pig4cloud.pigx.knowledge.service.AiAssistantService;
import com.pig4cloud.pigx.knowledge.service.AiDashscopeAssistantService;
import com.pig4cloud.pigx.knowledge.service.AiSiliconflowAssistantService;
import com.pig4cloud.pigx.knowledge.support.constant.ModelSupportEnums;
import com.pig4cloud.pigx.knowledge.support.provider.ModelProvider;
import com.pig4cloud.pigx.knowledge.support.util.PromptBuilder;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.Map;

/**
 * 文生图
 *
 * @author lengleng
 * @date 2024/3/26
 */
@Slf4j
@Component("text2ImageChat")
@RequiredArgsConstructor
public class Text2ImageChatRule implements ChatRule {

	private final AiDashscopeAssistantService aiDashscopeAssistantService;

	private final ModelProvider modelProvider;

	private final AiModelMapper aiModelMapper;

	/**
	 * 1. 调用Chat模型优化用户提交的提示词 2. 调用Image模型生成图片
	 * @param chatMessageDTO 聊天消息 dto
	 * @return {@link Flux }<{@link AiMessageResultDTO }>
	 */
	@SneakyThrows
	public Flux<AiMessageResultDTO> process(ChatMessageDTO chatMessageDTO) {
		// 调用Chat模型优化用户提交的提示词
		AiAssistantService assistantService = modelProvider.getAiAssistant().getValue();
		String content = assistantService.chat(chatMessageDTO.getConversationId(),
				PromptBuilder.render("gen-image.st", Map.of("input", chatMessageDTO.getContent())));

		AiModelEntity aiModelEntity = aiModelMapper
			.selectOne(Wrappers.<AiModelEntity>lambdaQuery().eq(AiModelEntity::getName, chatMessageDTO.getModelName()));

		if (aiModelEntity == null) {
			return Flux.just(new AiMessageResultDTO("图片模型不存在"));
		}

		// 调用硅基流动模型生成图片
		if (aiModelEntity.getProvider().equals(ModelSupportEnums.SILICONFLOW_FLUX_1_SCHNELL.getProvider())) {
			AiSiliconflowAssistantService imageAssistant = modelProvider
				.getFileAssistant(chatMessageDTO.getModelName());
			SiliconflowImageModelDTO.ImageModelRequst requst = new SiliconflowImageModelDTO.ImageModelRequst();
			requst.setModel(aiModelEntity.getModelName());
			requst.setPrompt(content + PromptBuilder.render("gen-image-limit.st", Map.of()));

			SiliconflowImageModelDTO.ImageModelResponse imageResponse = imageAssistant.generateImage(requst);
			String url = imageResponse.getImages().get(0).getUrl();
			// 转成 AiMessageResultDTO
			return Flux.just(new AiMessageResultDTO(PromptBuilder.render("gen-image-result.st", Map.of("url", url))));
		}

		// 调用dashscope模型生成图片
		String url = aiDashscopeAssistantService.generateImage(content, aiModelEntity);
		if (StrUtil.isNotBlank(url)) {
			return Flux.just(new AiMessageResultDTO(PromptBuilder.render("gen-image-result.st", Map.of("url", url))));
		}
		return Flux.just(new AiMessageResultDTO("生成图片失败"));

	}

}
